clearvars; clc; close all;

global option k_input beta mu_z sig_z c pro z0 rho lambda alp m_bar eta
%% parameters
z_min = -10;
z_max = 5;
I = 10e4; % number of grids

z = linspace(z_min,z_max,I)';
dz = z(2) - z(1);
dz2 = dz^2;

i_entry = round((z0 - z_min)/dz)+1;
z0 = z(i_entry);
psi = zeros(I,1);
psi(i_entry) = 1/dz;

% q_bar = [10; 50]; % regulatory thresholds: billion dollars 
q_bar = [10; 50; 250];
q_bar = log(q_bar);

switch option
    case 'fixed'
        %k = [0.0005; 0.0007];
        %k = [0.0005; 0.0007];
        k = k_input;
    case 'proportional'
        %k = [0.0163; 0.0046];
        %k = [0.0163; 0.0046];
        %k = [0.006;0.0045];
        %k = [0.0032; 0.0009];
        k = k_input;
    otherwise
        warning('Unexpected option')
end

kcum = cumsum(k);

R = 0.07; % initial guess of interest rate: mean of loan rate from 2001Q3 to 2007Q3

n_bar = length(q_bar); % number of thresholds
zh = zeros(n_bar, 1);
ih = zeros(I,n_bar);
il = zeros(I,n_bar);
io = zeros(I,n_bar);

%% iterations
iter = 0;
err = 1;
crit = 1e-6;
maxit = 1e4;
the = 0.01; % relaxation parameter on interest rate loop

while (err>crit) && (iter<=maxit)
    %% piecewise profit functions
    zl = q_bar - beta*R + 1;
    for j = 1:n_bar
        switch option
            case 'fixed'
                [zh(j),term] = broyden('z_high1',zl(j)+0.1,beta,k(j),R,q_bar(j)); %using k instead of kcum because the additional costs matters for bunching
            case 'proportional'
                [zh(j),term] = broyden('z_high2',zl(j)+0.1,beta,k(j),R,q_bar(j));
        end

        ih(:,j) = z - zh(j)>0; %indicate z is above z overline
        il(:,j) = z - zl(j)<0;
        io(:,j) = 1 - ih(:,j) - il(:,j);
    end
    if min(zh-zl)<=0
        warning('zh is lower than zl')
    end
    qstar = io*q_bar + (z + beta*R - 1).*(1 - sum(io,2)); % optimal size choice
    switch option
        case 'fixed'
            piz = (R - (qstar - z)/beta).*exp(qstar) - ih*kcum;
        case 'proportional'
            piz = (R - (qstar - z)/beta).*exp(qstar).*(1- ih*kcum); % profits, using kcum instead of k because the total costs matters for profits
    end
            
    %% HJB
    mu = mu_z*ones(I,1);
    sig = sig_z*ones(I,1);
    sig2 = sig.^2;
    
    X = -min(mu,0)/dz + sig2/(2*dz2);
    Y = -max(mu,0)/dz + min(mu,0)/dz - sig2/dz2;
    Z = max(mu,0)/dz + sig2/(2*dz2);

    A = spdiags(Y,0,I,I) + spdiags(X(2:I),-1,I,I) + spdiags([0;Z(1:I-1)],1,I,I);
    
    A(1,1) = Y(1) + X(1); % Reflecting barrier v'(z_min) = 0
    A(I,I) = Y(I) + Z(I); % Reflecting barrier v'(z_max) = 0
    
    B = (rho+lambda)*speye(I) - A;
    v = B\piz; % value functions derived from HJB
    
    %% KFE
    % entry
    m = m_bar*exp(eta*(sum(v.*psi.*dz)-c));

    Aa = A';
%     Aa(1,1) = 0; % Reflecting barrier f(z_min) = 0
%     Aa(I,I) = 0; % Reflecting barrier f(z_max) = 0
    Bb = -(Aa-lambda*speye(I)); % p.36-37, lecture 5-6 
    f = Bb\(m.*psi); % stationary distribution of z derived from KFE
    %f = f/sum(f);
    K = max(sum(exp(qstar).*f.*dz),1e-1);
    
    RR = pro*alp*K^(alp - 1); % lending market clearing condition
    R = (1 - the)*R + the*RR;
    err = abs(RR-R);
    iter = iter + 1;
    fprintf('Iteration number = %4d; R = %.6f; Error = %3.0e \n',iter,R,err);
end

%disp(zh)
%disp(zl)

%% figures
figure()
set(gca,'FontSize',12)
plot(z,v,'LineWidth',2)
grid
xlabel('z')
ylabel('$v\left(z\right)$','interpreter','latex')
xlim([z_min z_max])

figure()
set(gca,'FontSize',12)
plot(z,f,'LineWidth',2)
grid
xlabel('z')
ylabel('$f\left(z\right)$','interpreter','latex')
xlim([z_min z_max])
ylim([0 max(f)])

figure()
set(gca,'FontSize',12)
plot(z,qstar,'LineWidth',2)
%plot(z,exp(qstar),'LineWidth',2)
grid
xlabel('z')
ylabel('$q\left(z\right)$','interpreter','latex')
xlim([z_min z_max])

figure()
set(gca,'FontSize',12)
plot(z,piz./(exp(qstar)),'LineWidth',2)
grid
xlabel('z')
ylabel('Profit margin','interpreter','latex')
xlim([z_min z_max])



q0 = z + beta*R - 1;
q_distr = f.*(1 - sum(io,2));
for i = 1:n_bar
    bunch_l(i) = find(il(:,i),1,'last') ;
    if max(ih(:,i))==1
    bunch_h(i) = find(ih(:,i),1,'first') - 1;
    end
    q_distr(bunch_l) = sum(f(bunch_l:bunch_h));
end

ng = 25; % number of groups
ne = I/ng; 
q_dtr = reshape(q_distr,ne,ng);
q_dtr = sum(q_dtr);
q_dtr = repmat(q_dtr,ne,1);
q_dtr = reshape(q_dtr,I,1);

figure()
set(gca,'FontSize',12)
plot(exp(q0),q_dtr,'LineWidth',2)
grid
xlabel('Assets')
xlim([0 60])
ylabel('Frequency')

%disp(['size of entrants: ', num2str(exp(z0+beta*R-1))])
%disp(['average bank size: ', num2str(sum(exp(qstar).*f.*dz)/sum(f.*dz))])
%disp(['mass: ',num2str(sum(f.*dz))])
%disp(['Upper bound of bank size: ', num2str(max(exp(qstar)))])
%disp(['Lower bound of bank size: ' ,num2str(min(exp(qstar)))])

q0 = (z + beta*R - 1); %undistorted size
mb = v./(.1*exp(qstar)); % assuming book equity is 1/10 of assets


high = find(exp(q0)>50,1,'first');
median = find(exp(q0)>10,1,'first');
per_avoid = (sum(f(bunch_l(1):bunch_h(1)))+sum(f(bunch_l(2):bunch_h(2))))/sum(f); % percentage of banks avoiding regulation

mean_l = sum(f(1:median-1).*exp(qstar(1:median-1)).*dz)/sum(f(1:median-1).*dz);
mean_h = sum(f(high:end).*exp(qstar(high:end)).*dz)/sum(f(high:end).*dz);
mean_m = sum(f(median:high-1).*exp(qstar(median:high-1)).*dz)/sum(f(median:high-1).*dz);

asset_share_l = sum(f(1:median-1).*exp(qstar(1:median-1)).*dz)/sum(f.*exp(qstar).*dz);
asset_share_h = sum(f(high:end).*exp(qstar(high:end)).*dz)/sum(f.*exp(qstar).*dz);
asset_share_m = sum(f(median:high-1).*exp(qstar(median:high-1)).*dz)/sum(f.*exp(qstar).*dz);

num_share_l = sum(f(1:median-1).*dz)/sum(f.*dz);
num_share_h = sum(f(high:end).*dz)/sum(f.*dz);
num_share_m = sum(f(median:high-1).*dz)/sum(f.*dz);


pi_l = sum(f(1:median-1).*piz(1:median-1))/sum(f(1:median-1));
pi_m = sum(f(median:high-1).*piz(median:high-1))/sum(f(median:high-1));
pi_h = sum(f(high:end).*piz(high:end))/sum(f(high:end));

%mb_all = sum(v.*f.*dz)/sum(.1*exp(qstar).*f.*dz); % weighted average
mb_all = sum((v+exp(qstar)).*f.*dz)/sum(exp(qstar).*f.*dz); % weighted average

mb_l = sum(v(1:median-1).*f(1:median-1).*dz)/sum(.1*exp(qstar(1:median-1)).*f(1:median-1).*dz); % weighted average
mb_m = sum(v(median:high-1).*f(median:high-1).*dz)/sum(.1*exp(qstar(median:high-1)).*f(median:high-1).*dz); % weighted average
mb_h = sum(v(high:end).*f(high:end).*dz)/sum(.1*exp(qstar(high:end)).*f(high:end).*dz); % weighted average


% mb_l = sum(f(1:median-1).*mb(1:median-1))/sum(f(1:median-1));
% mb_m = sum(f(median:high-1).*mb(median:high-1))/sum(f(median:high-1));
% mb_h = sum(f(high:end).*mb(high:end))/sum(f(high:end));

welf_br = pro*K^alp;
welf_dp = sum((f.*dz.*(qstar).^2)/(2*beta));
welf_bk = sum(f.*piz)/sum(f);

mass = sum(f.*dz);


